'''
@IDE ：PyCharm
@Author ：知了-海龙
@Date ：2021/8/30 13:27
@Email：hallen200806@163.com
@Desc ：
'''

import os
import zipfile
from models.user.user_model import Docs
from sqlalchemy.orm import Session


def zip_file(output_file,file:Docs):
    zp = zipfile.ZipFile(output_file, "w")
    zp.write(file.filepath,file.name)
    zp.close()


def zip_dir(output_file,dir):
    zp = zipfile.ZipFile(output_file, "w")

    for path,dirname,filenames in os.walk(dir):
        fpath = path.replace(dir,"")
        for filename in filenames:
            zp.write(os.path.join(path,filename),os.path.join(fpath,filename))

    zp.close()

def zip_file_or_dir(db:Session,output_file,docs:[Docs]):
    zp = zipfile.ZipFile(output_file, "w")

    for doc in docs:
        if doc.flag == 1:  # 文件夹
            for path, dirname, filenames in os.walk(doc.filepath):
                fpath = path.replace(doc.filepath.rsplit("/",1)[0], "")
                for filename in filenames:
                    zp.write(os.path.join(path, filename), os.path.join(fpath, filename))

        else:
            p_doc = db.query(Docs).filter(Docs.id == doc.pid).first()
            fpath = doc.filepath.replace(p_doc.filepath, "")
            zp.write(doc.filepath, os.path.join(fpath, doc.name))

    zp.close()
